// Stdlib header file for input and output.
#include <iostream>

// Header file to access Pythia 8 program elements.
#include "Pythia.h"

// ROOT, for histogramming.
#include "TH1F.h"
#include "TH2F.h"
#include "TNtuple.h"

// ROOT, for saving file.
#include "TFile.h"

using namespace Pythia8;

int main(int argc, char* argv[]) {
  //Set number of events
  long nEvents;
  if(argc > 1){
    nEvents = atoi(argv[1]);
  }
  else{
    nEvents = 1000;
  }
  // Create Pythia instance and read settings from file
  Pythia pythia;
  pythia.readFile("task.cmnd");
  pythia.init();

  // Create file on which data can be saved.
  TFile* outFile = new TFile("task.root", "RECREATE");

  // Create histograms for overall data
  const Double_t Nmin=-0.5, Nmax=199.5, Nmax1=49.5, Nmax2=79.5, Nmax3=119.5;
  const Int_t Nbins=200, Nbins1=50, Nbins2=80, Nbins3=120;

  TNtuple particle_data ("particle_data","Various data per particle container", "id:eta:y:pT:pZ");
  TH1F *eta = new TH1F("eta","#eta distribution", 100, -10.0, 10.0);
  TH1F *y = new TH1F("y","y distribution", 100, -10.0, 10.0);
  TH1F *mult = new TH1F("mult","multiplicity", Nbins, Nmin, Nmax);

  //Create 2*3 multiplicity distibution histograms for 0.5, 1.0, 2.0 windows
  TH1F *Neta1 = new TH1F("Neta1", "multiplicity in eta<0.5 window", Nbins1, Nmin, Nmax1);
  TH1F *Ny1 = new TH1F("Ny1", "multiplicity in y<0.5 window", Nbins1, Nmin, Nmax1);
  TH1F *Neta2 = new TH1F("Neta2", "multiplicity in eta<1.0 window", Nbins2, Nmin, Nmax2);
  TH1F *Ny2 = new TH1F("Ny2", "multiplicity in y<1.0 window", Nbins2, Nmin, Nmax2);
  TH1F *Neta3 = new TH1F("Neta3", "multiplicity in eta<2.0 window", Nbins3, Nmin, Nmax3);
  TH1F *Ny3 = new TH1F("Ny3", "multiplicity in y<2.0 window", Nbins3, Nmin, Nmax3);

  //Create 2*3 2D Pt_N histograms for 0.5, 1.0, 2.0 windows
  const Double_t Ptmax=1.0;
  const Int_t Ptbins=75;
  TH2F *PtNeta1 = new TH2F("PtNeta1","<pT>:N in eta<0.5 window hist", Nbins1, Nmin, Nmax1, Ptbins, 0, Ptmax);
  TH2F *PtNy1 = new TH2F("PtNy1","<pT>:N in y<0.5 window hist", Nbins1, Nmin, Nmax1, Ptbins, 0, Ptmax);
  TH2F *PtNeta2 = new TH2F("PtNeta2","<pT>:N in eta<1 window hist", Nbins2, Nmin, Nmax2, Ptbins, 0, Ptmax);
  TH2F *PtNy2 = new TH2F("PtNy2","<pT>:N in y<1 window hist", Nbins2, Nmin, Nmax2, Ptbins, 0, Ptmax);
  TH2F *PtNeta3 = new TH2F("PtNeta3","<pT>:N in eta<2 window hist", Nbins3, Nmin, Nmax3, Ptbins, 0, Ptmax);
  TH2F *PtNy3 = new TH2F("PtNy3","<pT>:N in y<2 window hist", Nbins3, Nmin, Nmax3, Ptbins, 0, Ptmax);

  //Create 2 2D Pt_N histograms for 0.8 window (for ALICE data)
  TH2F *PtNetaexp = new TH2F("PtNetaexp","<pT>:N in eta<0.8 window hist", Nbins2, Nmin, Nmax2, Ptbins, 0, Ptmax);
  TH2F *PtNyexp = new TH2F("PtNyexp","<pT>:N in y<0.8 window hist", Nbins2, Nmin, Nmax2, Ptbins, 0, Ptmax);

  // Begin event loop. Generate event; skip if generation aborted.
  for (int iEvent = 0; iEvent < nEvents; ++iEvent) {
      if (!pythia.next()) continue;
      int N=0; //total multiplicity
      float PT[8]={0,0,0,0,0,0,0,0}; //will sum up pT here
      float Nlocal[8]={0,0,0,0,0,0,0,0}; //then divide by times summed up
      for (int i = 0; i < pythia.event.size(); ++i)
      //Check if particle is charged and in final state
        if (pythia.event[i].isFinal() && pythia.event[i].isCharged()){
		//Check if particle belong to the right window
		if(pythia.event[i].eta()<0.5&&pythia.event[i].eta()>-0.5){
		PT[0]+=pythia.event[i].pT();
		Nlocal[0]++;}
		if(pythia.event[i].y()<0.5&&pythia.event[i].y()>-0.5){
		PT[1]+=pythia.event[i].pT();
		Nlocal[1]++;}
		if(pythia.event[i].eta()<1&&pythia.event[i].eta()>-1){
		PT[2]+=pythia.event[i].pT();
		Nlocal[2]++;}
		if(pythia.event[i].y()<1&&pythia.event[i].y()>-1){
		PT[3]+=pythia.event[i].pT();
		Nlocal[3]++;}
		if(pythia.event[i].eta()<2&&pythia.event[i].eta()>-2){
		PT[4]+=pythia.event[i].pT();
		Nlocal[4]++;}
		if(pythia.event[i].y()<2&&pythia.event[i].y()>-2){
		PT[5]+=pythia.event[i].pT();
		Nlocal[5]++;}
		if(pythia.event[i].eta()<0.8&&pythia.event[i].eta()>-0.8){
		PT[6]+=pythia.event[i].pT();
		Nlocal[6]++;}
		if(pythia.event[i].y()<0.8&&pythia.event[i].y()>-0.8){
		PT[7]+=pythia.event[i].pT();
		Nlocal[7]++;}
	N++;
        eta->Fill(pythia.event[i].eta());
	y->Fill(pythia.event[i].y());
	//data follows id:eta:y:pT:pZ
	particle_data.Fill(pythia.event[i].id(),
			   pythia.event[i].eta(),
			   pythia.event[i].y(),
			   pythia.event[i].pT(),
			   pythia.event[i].pz());
      }
  for(int i=0; i<8; i++){
    PT[i]=PT[i]/Nlocal[i]; //Getting mean of PT
    }
  PtNeta1->Fill(Nlocal[0],PT[0]);
  PtNy1->Fill(Nlocal[1],PT[1]);
  PtNeta2->Fill(Nlocal[2],PT[2]);
  PtNy2->Fill(Nlocal[3],PT[3]);
  PtNeta3->Fill(Nlocal[4],PT[4]);
  PtNy3->Fill(Nlocal[5],PT[5]);

  PtNetaexp->Fill(Nlocal[6],PT[6]);
  PtNyexp->Fill(Nlocal[7],PT[7]);

  mult->Fill(N);

  Neta1->Fill(Nlocal[0]);
  Ny1->Fill(Nlocal[1]);
  Neta2->Fill(Nlocal[2]);
  Ny2->Fill(Nlocal[3]);
  Neta3->Fill(Nlocal[4]);
  Ny3->Fill(Nlocal[5]);
  }

  // Save histogram on file and close file.
  mult->Write();
  eta->Write();
  y->Write();
  particle_data.Write();

  PtNeta1->Write();
  PtNy1->Write();
  PtNeta2->Write();
  PtNy2->Write();
  PtNeta3->Write();
  PtNy3->Write();

  PtNetaexp->Write();
  PtNyexp->Write();

  Neta1->Write();
  Ny1->Write();
  Neta2->Write();
  Ny2->Write();
  Neta3->Write();
  Ny3->Write();

  delete outFile;

  // Done.
  return 0;
}
